"""
10_multi_agent.py - Multi-Agent Orchestration
===============================================
Chapter 12: Building systems where multiple AI agents collaborate.

KEY CONCEPTS:
  ★ Supervisor Pattern - One agent delegates to specialists
  ★ Agent Loop - Supervisor routes, specialist works, return to supervisor
  ★ Specialist Agents - Domain-specific processing
  ★ Dynamic Routing - Supervisor decides next agent at runtime

Run: python src/10_multi_agent.py
"""

from typing import TypedDict, Annotated, Literal
from operator import add
from langgraph.graph import StateGraph, START, END


class TeamState(TypedDict):
    """State for the multi-agent team."""
    task: str
    messages: Annotated[list, add]    # Accumulates all messages
    next_agent: str
    result: str
    iterations: int


# ============================================================
# SPECIALIST AGENTS
# ============================================================

def researcher(state: TeamState) -> dict:
    """Research specialist - gathers information."""
    task = state["task"]
    return {
        "messages": [f"[Researcher] Gathered data on: {task}"],
        "result": f"Research findings on {task}",
        "iterations": state.get("iterations", 0) + 1
    }


def writer(state: TeamState) -> dict:
    """Writing specialist - creates content."""
    return {
        "messages": [f"[Writer] Drafted content based on research"],
        "result": "Written draft complete",
        "iterations": state.get("iterations", 0) + 1
    }


def reviewer(state: TeamState) -> dict:
    """Review specialist - quality checks."""
    return {
        "messages": [f"[Reviewer] Quality review complete. Approved."],
        "result": "Review passed",
        "iterations": state.get("iterations", 0) + 1
    }


# ============================================================
# SUPERVISOR AGENT
# ============================================================

def supervisor(state: TeamState) -> dict:
    """Supervisor decides which agent to call next.
    
    ★ In production, this would use an LLM to decide.
    Here we use a simple sequential workflow:
    research -> write -> review -> FINISH
    """
    iterations = state.get("iterations", 0)
    
    if iterations == 0:
        return {"next_agent": "researcher", "messages": ["[Supervisor] Starting with research"]}
    elif iterations == 1:
        return {"next_agent": "writer", "messages": ["[Supervisor] Moving to writing"]}
    elif iterations == 2:
        return {"next_agent": "reviewer", "messages": ["[Supervisor] Final review"]}
    else:
        return {"next_agent": "FINISH", "messages": ["[Supervisor] All done!"]}


def route_to_agent(state: TeamState) -> str:
    """Route based on supervisor's decision."""
    agent = state["next_agent"]
    if agent == "FINISH":
        return "__end__"
    return agent


# ============================================================
# BUILD MULTI-AGENT GRAPH
# ============================================================

graph = StateGraph(TeamState)

graph.add_node("supervisor", supervisor)
graph.add_node("researcher", researcher)
graph.add_node("writer", writer)
graph.add_node("reviewer", reviewer)

graph.add_edge(START, "supervisor")
graph.add_conditional_edges("supervisor", route_to_agent)

# After each specialist, return to supervisor
graph.add_edge("researcher", "supervisor")
graph.add_edge("writer", "supervisor")
graph.add_edge("reviewer", "supervisor")

app = graph.compile()


if __name__ == "__main__":
    print("=" * 60)
    print("10: Multi-Agent Orchestration (Supervisor Pattern)")
    print("=" * 60)
    
    result = app.invoke({"task": "Write a blog post about AI agents"})
    
    print("\n📋 Execution Trace:")
    for msg in result["messages"]:
        print(f"  {msg}")
    
    print(f"\n  Final Result: {result['result']}")
    print(f"  Total Iterations: {result['iterations']}")
    
    print(f"\n{'=' * 60}")
    print("✅ Multi-agent orchestration complete!")
    print("=" * 60)
